import websocket  # NOTE: websocket-client (https://github.com/websocket-client/websocket-client)
import uuid
import json
import random
import sys
from sys import argv
import urllib.request
import urllib.parse

# 本地ip 默认127.0.0.1:8188
server_address = "192.168.1.36:8188"
# 输出路径
savePath = "G:\\ComfyUI-aki\\temp\\"
client_id = str(uuid.uuid4())
# 区分文本
name = sys.argv[0]
print("开始运行", name)


def queue_prompt(prompt):
    p = {"prompt": prompt, "client_id": client_id}
    data = json.dumps(p).encode('utf-8')
    req = urllib.request.Request("http://{}/prompt".format(server_address), data=data)
    return json.loads(urllib.request.urlopen(req).read())


def get_image(filename, subfolder, folder_type):
    data = {"filename": filename, "subfolder": subfolder, "type": folder_type}
    url_values = urllib.parse.urlencode(data)
    with urllib.request.urlopen("http://{}/view?{}".format(server_address, url_values)) as response:
        return response.read()


def get_history(prompt_id):
    with urllib.request.urlopen("http://{}/history/{}".format(server_address, prompt_id)) as response:
        return json.loads(response.read())


def get_images(ws, prompt):
    prompt_id = queue_prompt(prompt)['prompt_id']
    output_images = {}
    while True:
        out = ws.recv()
        if isinstance(out, str):
            message = json.loads(out)
            print(message)
            if message['type'] == 'executing':
                data = message['data']
                print(data)
                if data['node'] is None and data['prompt_id'] == prompt_id:
                    break  # Execution is done
        else:
            # If you want to be able to decode the binary stream for latent previews, here is how you can do it:
            # bytesIO = BytesIO(out[8:])
            # preview_image = Image.open(bytesIO) # This is your preview in PIL image format, store it in a global
            continue  # previews are binary data

    history = get_history(prompt_id)[prompt_id]
    if not history['outputs']:
        for messages in history["status"]:
            message = history["status"]["messages"][2]
            return [history["status"][messages], message[1]["exception_message"], message[1]["exception_type"]]
    else:
        for node_id in history['outputs']:
            node_output = history['outputs'][node_id]
            images_output = []
            imageName = ""
            if 'images' in node_output:
                for image in node_output['images']:
                    imageName = image['filename']
                    image_data = get_image(image['filename'], image['subfolder'], image['type'])
                    images_output.append(image_data)

            output_images[node_id] = images_output

        return ["success", output_images, imageName]

random_num = random.randint(1, 2147483647)
# 输入json工作流
# file_path = "C:\\\\Users\\\\Administrator\\\\Downloads\\\\flux+lora.json"
file_path = "\\\\Nas-02\\\\常用软件\\\\研发一部\\\\研发一部\\\\model\\\\flux+lora.json"
with open(file_path, 'r', encoding='utf-8') as file:
    prompt = json.loads(file.read())
    # 调节工作流参数
    prompt["25"]["inputs"]["noise_seed"] = random_num
    prompt["17"]["inputs"]["steps"] = 1
# ws监听进度并输出图片信息
ws = websocket.WebSocket()
ws.connect("ws://{}/ws?clientId={}".format(server_address, client_id))
ws.send("Hello, Server")
imageArr = get_images(ws, prompt)
images = imageArr[1]
history = imageArr[2]
ws.close()
if "success" in imageArr:
    print("运行完成", name, savePath + history)
else:
    print("运行失败", name, images + history)

